#!/bin/bash
# 完整测试流程脚本

set -e

echo "=========================================="
echo "Full Testing Pipeline"
echo "=========================================="

# 配置
NUM_GPUS=${NUM_GPUS:-1}
SKIP_DOWNLOAD=${SKIP_DOWNLOAD:-0}

echo "Configuration:"
echo "  NUM_GPUS: $NUM_GPUS"
echo "  SKIP_DOWNLOAD: $SKIP_DOWNLOAD"
echo "=========================================="

# Step 1: 下载模型和数据（如果需要）
if [ "$SKIP_DOWNLOAD" -eq 0 ]; then
    echo ""
    echo "Step 1: Downloading models and datasets..."
    cd tools && python download_models.py --model gpt2-medium --data && cd ..
fi

# Step 2: PyTorch基线训练
echo ""
echo "Step 2: Running PyTorch baseline training..."
cd train
if [ "$NUM_GPUS" -gt 1 ]; then
    NUM_GPUS=$NUM_GPUS bash run_pytorch.sh
else
    bash run_pytorch.sh
fi
cd ..

# Step 3: Triton优化训练
echo ""
echo "Step 3: Running Triton optimized training..."
cd train
if [ "$NUM_GPUS" -gt 1 ]; then
    NUM_GPUS=$NUM_GPUS bash run_triton.sh
else
    bash run_triton.sh
fi
cd ..

# Step 4: 性能对比
echo ""
echo "Step 4: Comparing results..."
cd tools && python compare_results.py && cd ..

echo ""
echo "=========================================="
echo "Full testing pipeline completed!"
echo "=========================================="
echo ""
echo "Results available in:"
echo "  - PyTorch: results/pytorch_baseline/"
echo "  - Triton: results/triton_optimized/"
echo "  - Comparison: results/comparison/"
echo ""

